package edu.vt.cbil.core;

import edu.vt.cbil.dataset.BioCollection;
import edu.vt.cbil.dataset.BioSample;
import edu.vt.cbil.dataset.ProbeSet;
import edu.vt.cbil.dataset.Segment;
import edu.vt.cbil.util.BioToolbox;
import org.apache.commons.math3.analysis.function.Gaussian;

import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.logging.Logger;

/**
 * The type CN segment cutter.
 *
 * @author Justin
 * @version 1.0
 * @created 12 -11月-2014 21:40:04
 */
public class CNSegmentCutter {
    private double PVALUE_THRESH = 0.05;
    private int MIN_SEG_LENGTH = 256;
    private int MIN_SEG_STEP = 8;
    private double COMBINE_FACTOR = 0.2;
    private Logger m_log;
    private BioCollection p_Collection;
    private BioSample p_sample;
    private ArrayList<Segment> m_segments;
    private double[] weightFactor;

    /**
     * Instantiates a new CN segment cutter.
     *
     * @param pc the BioCollection
     */
    public CNSegmentCutter(BioCollection pc) {

        m_log = Logger.getLogger("edu.vt.cbil.bacom2");

        p_Collection = pc;

        m_segments = new ArrayList<Segment>();

        for (int i = 0; i < 27; i++) {
            m_segments.add(new Segment());
        }

    }

    public CNSegmentCutter(BioCollection pc, double pvalueThre, int minSegLen, int minStep, double combinefactor) {

        m_log = Logger.getLogger("edu.vt.cbil.bacom2");

        p_Collection = pc;

        m_segments = new ArrayList<Segment>();

        for (int i = 0; i < 27; i++) {
            m_segments.add(new Segment());
        }

        this.PVALUE_THRESH = pvalueThre;
        this.MIN_SEG_LENGTH = minSegLen;
        this.MIN_SEG_STEP = minStep;
        this.COMBINE_FACTOR = combinefactor;

    }

    public void setSample(BioSample sample) {
        this.p_sample = sample;
    }

    /**
     * Sets collection.
     *
     * @param pc the pc
     */
    public void setCollection(BioCollection pc) {
        this.p_Collection = pc;
    }

    /**
     * Split chromosomes to segments.
     *
     * @param p_sample the p _ sample
     */
    public void splitChromosomes(BioSample p_sample) {

        m_log.info("Parallel computing starting ... ...");

        int threadNum = p_Collection.ChrIdInterested.size();
        CountDownLatch threadSignal = new CountDownLatch(threadNum);

        for (Short i : p_Collection.ChrIdInterested) {
            Segment input = m_segments.get(i);
            Set<Segment> output = p_sample.getAllSegments(i);

            Thread t = new Cutter(input, output, threadSignal);
            t.start();
        }

        try {
            threadSignal.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        combineSegments(p_sample, COMBINE_FACTOR);
    }

    /**
     * Prepare copy number segment.
     *
     * @param p_sample the p _ sample
     */
    public void prepareCopyNumberSegment(BioSample p_sample) {

        for (Short i : p_Collection.getChrIdArray()) {
            Segment seg = m_segments.get(i);
            seg.Loci.clear();

            seg.Chr_id = i;
            seg.Seg_id = 0L;

            for (ProbeSet ps : p_sample.getAllProbeSets(i).values()) {
                if (ps.data.Is_Valid) {
                    seg.Loci.put(ps.Loci, ps.data.CopyNumberHalf);
                    if (seg.Seg_id == 0) {
                        seg.StartLoci = ps.Loci;
                        seg.Seg_id = seg.StartLoci;
                    }
                }
            }

            seg.resetIntegralCN();
            seg.resetIndex();
            seg.updateCharacters(PVALUE_THRESH);

            m_log.info(seg.getCharacterString());
        }
    }

    /**
     * Smoothing the copynumber array.
     *
     * @param w the width of smoothing windows
     */
    public void smoothing(int w) {

        int width = w | 0x01;
        int halfwidth = width / 2;

        updateWeightFactor(halfwidth);

        for (Segment seg : m_segments) {
            int i = 0;
            int len = seg.Loci.size();
            if (len > 0) {
                ArrayList<Double> arr = new ArrayList<Double>(seg.Loci.values());

                for (Map.Entry<Long, Double> d : seg.Loci.entrySet()) {

                    if ((i >= halfwidth) && (i < len - halfwidth)) {

                        double nv = addWeight(arr, i - halfwidth, i + halfwidth);

                        d.setValue(nv);
                    }
                    i++;

                }
            }

        }
    }

    private double addWeight(ArrayList<Double> input, int start, int stop) {

        double sum = 0;

        for (int i = start; i <= stop; i++) {
            sum += input.get(i) * weightFactor[i - start];
        }

        return sum;

    }

    private void updateWeightFactor(int half) {

        weightFactor = new double[2 * half + 1];

        Gaussian g = new Gaussian(0, 1);
        double step = 3.0 / half;
        double pos = step;
        double sum;

        weightFactor[half] = g.value(0);
        sum = weightFactor[half];
        for (int i = 1; i <= half; i++, pos += step) {
            double v = g.value(pos);
            weightFactor[half + i] = v;
            weightFactor[half - i] = v;
            sum += 2 * v;
        }
        for (int i = 0; i < weightFactor.length; i++) {
            weightFactor[i] = weightFactor[i] / sum;
        }

    }


    /**
     * Calc segment mean _ std.
     *
     * @param input the input
     */
    public void calcSegmentMean_Std(Segment input) {

        input.updateCharacters(PVALUE_THRESH);
        System.out.println(input.getCharacterString());
    }

    /**
     * Combine segments by the robust mean of each segment.
     *
     * @param p_sample the p _ sample
     */
    public void combineSegments(BioSample p_sample, double combineFactor) {

        List<Double> diffList = new LinkedList<Double>();
        for (Short i : p_Collection.getChrIdArray()) {
            Set<Segment> segments = p_sample.getAllSegments(i);

            boolean notFirst = false;
            double last = 0;

            for (Segment s : segments) {
                double value = BioToolbox.log2(s.robustMeanHalfCopyNumber);
                if (notFirst) {
                    diffList.add(Math.abs(value - last));
                } else {
                    notFirst = true;
                }
                last = value;
            }
        }
        Collections.sort(diffList);
        double std = BioToolbox.std(diffList.subList(0, diffList.size() >> 1));
        double std2 = BioToolbox.std(diffList);
        m_log.info(String.format("std of Mean HalfCopyNumber of segments = %.4f:%.4f; seg count=%d",
                std, std2, diffList.size() + 1));
        double f = std2 / std;


        double thresh = Math.max(0.03, diffList.get((int) (combineFactor * diffList.size())));

        for (Short i : p_Collection.getChrIdArray()) {
            Set<Segment> segments = p_sample.getAllSegments(i);

            segments = combineSegmentsInOneChromosome(segments, thresh, 0.2);
            segments = combineSegmentsInOneChromosome(segments, 0.05 * combineFactor, 0.2);

            p_sample.modifySegments(i, segments);
        }


    }

    private Set<Segment> combineSegmentsInOneChromosome(Set<Segment> segments, double avgLimit, double stdLimit) {

        int segSize = segments.size();
        double[] cn = new double[segSize];
        double[] stdcn = new double[segSize];
        double[] segLen = new double[segSize];
        int k = 0;
        for (Segment s : segments) {
            cn[k] = BioToolbox.log2(s.robustMeanHalfCopyNumber);
            stdcn[k] = s.robustStdHalfCopyNumber;
            segLen[k] = s.Loci.size();
            k++;
        }
        double[] fcn;

        if (cn.length > 10) {
            fcn = BioToolbox.GaussianBlur(cn, 5, 1);
        } else {
            fcn = cn;
        }

        List<Integer> breakpoints = new LinkedList<Integer>();

        double[] diffcn = new double[fcn.length];
        diffcn[0] = 0;
        //breakpoints.add(0);
        for (int j = 1; j < fcn.length; j++) {
            diffcn[j] = Math.abs(fcn[j] - fcn[j - 1]);
        }

//        double[] sortDiff = Arrays.copyOf(diffcn, diffcn.length);
//        Arrays.sort(sortDiff);
//
//        double avgLimit = Math.max(0.03,sortDiff[(int)(0.90*sortDiff.length)]);

        for (int j = 1; j < fcn.length; j++) {
            //if ((diffcn[j] < 0.05)||(segLen[j] < MIN_SEG_LENGTH))
            //这里的参数从严设置效果比较好一点
            if (((diffcn[j] < avgLimit)) && (stdcn[j] < stdLimit) && (stdcn[j - 1] < stdLimit)) {
                diffcn[j] = 0;
            } else {
                breakpoints.add(j);
            }
        }
        breakpoints.add(segSize);

        Set<Segment> newSegs = new TreeSet<Segment>();
        Iterator<Integer> it = breakpoints.iterator();
        Iterator<Segment> itSegs = segments.iterator();

        int begin;
        int end = 0;
        while (it.hasNext()) {
            begin = end;
            end = it.next();
            Segment curSeg = itSegs.next();
            for (int j = begin + 1; j < end; j++) {
                curSeg.Loci.putAll(itSegs.next().Loci);
            }
            newSegs.add(curSeg);
        }

        for (Segment ns : newSegs) {
            ns.updateCharacters();
        }
        return newSegs;
    }

    private class Cutter extends Thread {
        private Segment p_inputSegment;
        private Set<Segment> p_outputSegment;
        private CountDownLatch p_threadsSignal;
        private short Chr_id;

        /**
         * Instantiates a new Cutter.
         *
         * @param input         the input
         * @param output        the output
         * @param threadsSignal the threads signal
         */
        public Cutter(Segment input, Set<Segment> output, CountDownLatch threadsSignal) {

            p_inputSegment = input;
            p_outputSegment = output;
            this.p_threadsSignal = threadsSignal;
            Chr_id = input.Chr_id;
        }

        @Override
        public void run() {

            splitChromosome(p_inputSegment, p_outputSegment);

            m_log.info(String.format("#### #### #### chr %02d First Step: Loci Count = %05d; \t Segments Count = %d",
                    p_inputSegment.Chr_id, p_inputSegment.Loci.size(), p_outputSegment.size()));

            p_threadsSignal.countDown();
        }

        private void splitChromosome(Segment input, Set<Segment> output) {

            if (input.Loci.size() == 0) return;

            double maxZ = 0;
            int Start = 0;
            int Stop = 0;

            Queue<Segment> splittingSegments = new LinkedList<Segment>();
            Set<Segment> scrapSegments = new LinkedHashSet<Segment>();

            splittingSegments.add(input);
            List<Integer> Index = new ArrayList<Integer>();
            int newSegCount = 0;

            while (!splittingSegments.isEmpty()) {

                Segment curSeg = splittingSegments.poll();

                //̫如果太短
                if (curSeg.Loci.size() < 2 * MIN_SEG_LENGTH) {
                    scrapSegments.add(curSeg);
                } else {

                    maxZ = 0.0;
                    Index.clear();
                    int segLength = curSeg.Loci.size();

                    for (int i = 0; i < segLength; i += MIN_SEG_STEP) {
                        Index.add(i);
                    }
                    Index.set(Index.size() - 1, segLength - 1);

                    int indexLength = Index.size();
                    double zij;
                    int offset = MIN_SEG_LENGTH / MIN_SEG_STEP;
                    for (int i = 0; i < indexLength - offset; i++) {
                        for (int j = i + offset; j < indexLength - offset; j++) {
                            //如果过于接近整个段
                            //if ((i<MIN_SEG_LENGTH/MIN_SEG_STEP) || (j>indexLength - MIN_SEG_LENGTH/MIN_SEG_STEP))
                            if ((Index.get(j) - Index.get(i) + 1) >= segLength - MIN_SEG_LENGTH)
                                continue;

                            zij = curSeg.calculateZ(Index.get(i), Index.get(j));
                            if (zij > maxZ) {
                                maxZ = zij;
                                Start = Index.get(i);
                                Stop = Index.get(j);
                            }
                        }
                    }

                    if (maxZ > curSeg.splittingThreshold) {
                        //splitting
                        List<Segment> newSegs = new LinkedList<Segment>();

                        if ((Stop < Start) || (Stop == Start)) {//debug info
                            System.out.println("splitChromosome Error!");
                            System.exit(-1);
                        }

                        splitSegment(curSeg, newSegs, Start, Stop);

                        if (newSegs.size() > 0) {
                            splittingSegments.addAll(newSegs);
                        }
                    } else {
                        //add to result
                        curSeg.isReady = true;

                        output.add(curSeg);

                        newSegCount++;
                    }
                }


                if (splittingSegments.isEmpty()) {
                    if (newSegCount > 0) {
                        splittingSegments.addAll(mergeScrapSegments(scrapSegments, output));
                        scrapSegments.clear();
                        newSegCount = 0;
                    } else {
                        output.addAll(scrapSegments);
                        //output.addAll(mergeScrapSegments(scrapSegments,output));
                    }
                }

                //System.out.println(String.format("S: %d \t E: %d \t Z = %f", Start, Stop, maxZ));
            }
//            output.addAll(scrapSegments);
        }

        private List<Segment> mergeScrapSegments(Set<Segment> segments, Set<Segment> readySegments) {

            List<Segment> tempList = new LinkedList<Segment>();
            List<Segment> result = new LinkedList<Segment>();

            tempList.addAll(segments);
            tempList.addAll(readySegments);
            Collections.sort(tempList);

            Segment mergeSeg, nextSeg;
            Iterator<Segment> its = tempList.listIterator();

            boolean isMerging;
            while (its.hasNext()) {
                mergeSeg = its.next();
                isMerging = false;

                if (!mergeSeg.isReady) {
                    while (its.hasNext()) {
                        nextSeg = its.next();
                        if (!nextSeg.isReady) {
                            mergeSeg.Loci.putAll(nextSeg.Loci);
                            isMerging = true;
                        } else {
                            if (isMerging) {
                                result.add(mergeSeg);
                            } else {
                                readySegments.add(mergeSeg);
                            }
                            isMerging = false;
                            mergeSeg = null;
                            break;
                        }
                    }
                    if (mergeSeg != null) {
                        if (isMerging) {
                            result.add(mergeSeg);
                        } else {
                            readySegments.add(mergeSeg);
                        }
                    }
                }

            }
            for (Segment oneSeg : result) {
                oneSeg.resetIntegralCN();
                oneSeg.resetIndex();

                oneSeg.updateCharacters(PVALUE_THRESH);
            }

            return result;
        }

        private void splitSegment(Segment input, List<Segment> output, int start, int stop) {

            int s1 = (start > MIN_SEG_LENGTH) ? 1 : 0;
            int s2 = (stop < (input.Loci.size() - MIN_SEG_LENGTH)) ? 1 : 0;
            int s = (s1 << 1) | s2;

            switch (s) {
                case 0x00:
                case 0x03:
                    output.add(input.getSubSegment(0, start));
                    output.add(input.getSubSegment(start, stop));
                    output.add(input.getSubSegment(stop, input.Loci.size()));
                    break;
                case 0x02:
                    //cut the first part to new segment, if the third part is too short
                    output.add(input.getSubSegment(0, start));
                    output.add(input.getSubSegment(start, input.Loci.size()));
                    break;
                case 0x01:
                    //cut the third part to new segment, if the first part is too short
                    output.add(input.getSubSegment(0, stop));
                    output.add(input.getSubSegment(stop, input.Loci.size()));
                    break;
//                case 0x00:
//                    // if the first and third part all are too short.
//                    // do nothing
//                    break;
                default:

            }
            //update statistical Characters
            for (Segment oneSeg : output) {
                oneSeg.resetIntegralCN();
                oneSeg.resetIndex();


                oneSeg.updateCharacters(PVALUE_THRESH);

            }

        }
    }

}


